# -*- coding:utf-8 -*-
__author__ = 'Dragon Sun'
__date__ = '2025-07-02 19:18:46'

import datetime
import json
import time
from typing import Optional, Any, Callable

from websocket import WebSocketApp, WebSocket


class DSWebSocket(object):

    def __init__(self, url: str, ping_interval: int = 10):
        super().__init__()

        # 回调
        self.统计数据变化回调: Optional[Callable[[DSWebSocket], None]] = None
        self.发生错误回调: Optional[Callable[[DSWebSocket, Any], None]] = None

        self.url = url
        self.ping_interval = ping_interval  # 自动发送 ping 的时间间隔(秒) 由于库本身Bug，第一次发送Ping会是这个间隔的两倍

        self.手动连接时间: Optional[datetime.datetime] = None
        self.自动连接次数: int = 0
        self.最后一次连接时间: Optional[datetime.datetime] = None
        self.pong次数: int = 0
        self.最后一次pong时间: Optional[datetime.datetime] = None
        self.错误次数: int = 0
        self.最后一次错误时间: Optional[datetime.datetime] = None

        self.wss = WebSocketApp(
            url=url,
            keep_running=True,  # 是否在连接关闭后保持事件循环运行, 控制是否自动重连
            on_open=self._on_open,
            on_message=self._on_message,
            on_error=self._on_error,
            on_close=self._on_close,
            on_pong=self._on_pong,
        )

    def run(self, 自动重连: bool = False):
        if 自动重连:
            self.自动连接次数 += 1
        else:
            self.自动连接次数 = 0
            self.手动连接时间 = datetime.datetime.now()
        self._统计数据变化处理()

        self.wss.run_forever(
            ping_interval=self.ping_interval,  # 自动发送 ping 的时间间隔(秒), 20 (每20秒发送一次ping)
            ping_timeout=5,  # 等待 pong 响应的超时时间(秒), 检测死连接
            ping_payload='ping',  # ping 帧携带的自定义数据, 示例: b"heartbeat"
        )

    def _统计数据变化处理(self):
        if self.统计数据变化回调:
            self.统计数据变化回调(self)

    def _发生错误处理(self, 错误: Any):
        if self.发生错误回调:
            self.发生错误回调(self, 错误)

    def _on_open(self, ws: WebSocket):
        self.最后一次连接时间 = datetime.datetime.now()
        self._统计数据变化处理()
        ws.send('ping')

    def _on_message(self, ws: WebSocket, message: Any):
        pass

    def _on_error(self, ws: WebSocket, error: Any):
        self.错误次数 += 1
        self.最后一次错误时间 = datetime.datetime.now()
        self._统计数据变化处理()
        self._发生错误处理(error)
        self._on_close(ws, None, error)

    # def _on_close(self, ws: WebSocket, status_code: Any, reason: Any):
    def _on_close(self, _: WebSocket, __: Any, ___: Any):
        # 实现自动重连
        time.sleep(2)  # 等待2秒后重连
        self.run(自动重连=True)

    # def _on_pong(self, ws: WebSocket, data: Any):
    def _on_pong(self, _: WebSocket, __: Any):
        self.pong次数 += 1
        self.最后一次pong时间 = datetime.datetime.now()
        self._统计数据变化处理()


class DSOKXv5WebSocketBase(DSWebSocket):

    def __init__(self, url: str):
        # 由于库的Bug，第一次ping触发是ping_interval的两倍时间
        # 而OKXv5的WebSocket30秒无数据就会断开，所以ping_interval的值应该<30/2
        # 所以这里设置为10
        super().__init__(url=url, ping_interval=10)

    def 订阅(self, params: list[dict]):
        payload = json.dumps({
            "op": "subscribe",
            "args": params
        })
        self.wss.send(payload)
